//! K线生成器模块
//! 
//! 提供K线数据生成和处理功能

use chrono::{DateTime, Utc, Duration, Datelike};
use std::collections::HashMap;

use crate::objects::{bar::{NewBar, RawBar}, enums::Freq};

/// K线生成器
pub struct BarGenerator {
    /// 基础频率
    pub base_freq: String,
    /// 目标频率列表
    pub freqs: Vec<String>,
    /// 最大缓存数量
    pub max_count: usize,
    /// 市场类型
    pub market: String,
    /// 各频率的K线缓存
    pub freq_bars: HashMap<String, Vec<RawBar>>,
}

impl BarGenerator {
    /// 创建新的K线生成器
    pub fn new(base_freq: String, freqs: Vec<String>, max_count: usize, market: String) -> Self {
        let mut freq_bars = HashMap::new();
        for freq in &freqs {
            freq_bars.insert(freq.clone(), Vec::new());
        }

        Self {
            base_freq,
            freqs,
            max_count,
            market,
            freq_bars,
        }
    }

    /// 验证频率设置
    fn validate_freqs(&self) -> Result<(), String> {
        if !self.freqs.contains(&self.base_freq) {
            return Err(format!("base_freq {} 不在 freqs 列表中", self.base_freq));
        }

        // 检查频率顺序
        let freq_order = vec![
            "1分钟", "5分钟", "15分钟", "30分钟", "60分钟",
            "日线", "周线", "月线", "季线", "年线"
        ];

        for i in 0..self.freqs.len() - 1 {
            let current = &self.freqs[i];
            let next = &self.freqs[i + 1];
            
            let current_idx = freq_order.iter().position(|&x| x == current);
            let next_idx = freq_order.iter().position(|&x| x == next);
            
            if let (Some(current_pos), Some(next_pos)) = (current_idx, next_idx) {
                if current_pos >= next_pos {
                    return Err(format!("频率顺序错误: {} 应该在 {} 之前", current, next));
                }
            }
        }

        Ok(())
    }

    /// 初始化指定频率的K线
    pub fn init_freq_bars(&mut self, freq: &str, bars: Vec<RawBar>) -> Result<(), String> {
        if !self.freqs.contains(&freq.to_string()) {
            return Err(format!("频率 {} 不在支持的频率列表中", freq));
        }

        self.freq_bars.insert(freq.to_string(), bars);
        Ok(())
    }

    /// 更新指定频率的K线
    fn update_freq(&mut self, bar: &RawBar, freq: &Freq) -> Result<(), String> {
        let freq_str = freq.to_string();
        if !self.freqs.contains(&freq_str) {
            return Ok(()); // 不支持的频率直接跳过
        }

        // 作用域1：只做 bars 的可变借用
        {
            let bars = self.freq_bars.get_mut(&freq_str).unwrap();
            // 如果是第一个K线，直接添加并提前返回
            if bars.is_empty() {
                bars.push(bar.clone());
                return Ok(());
            }
        }

        // 作用域2：只做不可变借用
        let last_bar_clone = {
            let bars = self.freq_bars.get(&freq_str).unwrap();
            bars.last().unwrap().clone()
        };
        let need_new_bar = self.check_need_new_bar(bar, &last_bar_clone, freq);

        // 作用域3：准备写入的数据
        let (do_new, new_bar_opt);
        let (do_update, update_idx);
        if need_new_bar {
            new_bar_opt = Some(self.generate_new_bar(bar, freq)?);
            do_new = true;
            do_update = false;
            update_idx = 0;
        } else {
            new_bar_opt = None;
            do_new = false;
            do_update = true;
            let bars = self.freq_bars.get(&freq_str).unwrap();
            update_idx = bars.len() - 1;
        }

        // 作用域4：再次可变借用 bars，只做数据写入
        let bars = self.freq_bars.get_mut(&freq_str).unwrap();
        if do_new {
            bars.push(new_bar_opt.unwrap());
            if bars.len() > self.max_count {
                bars.remove(0);
            }
        } else if do_update {
            let last_bar = &mut bars[update_idx];
            // self.update_current_bar(last_bar, bar);
            last_bar.high = last_bar.high.max(bar.high);
            last_bar.low = last_bar.low.min(bar.low);
            last_bar.close = bar.close;
            last_bar.vol += bar.vol;
            last_bar.amount += bar.amount;
        }

        Ok(())
    }

    /// 检查是否需要生成新的K线
    fn check_need_new_bar(&self, bar: &RawBar, last_bar: &RawBar, freq: &Freq) -> bool {
        match freq {
            Freq::Min1 | Freq::Min5 | Freq::Min15 | Freq::Min30 | Freq::Min60 => {
                // 分钟级别，检查时间是否跨周期
                let freq_minutes = match freq {
                    Freq::Min1 => 1,
                    Freq::Min5 => 5,
                    Freq::Min15 => 15,
                    Freq::Min30 => 30,
                    Freq::Min60 => 60,
                    _ => return false,
                };
                
                let duration = Duration::minutes(freq_minutes);
                bar.dt >= last_bar.dt + duration
            }
            Freq::Day => {
                // 日线级别，检查是否跨日
                bar.dt.date_naive() != last_bar.dt.date_naive()
            }
            Freq::Week => {
                // 周线级别，检查是否跨周
                bar.dt.iso_week() != last_bar.dt.iso_week()
            }
            Freq::Month => {
                // 月线级别，检查是否跨月
                bar.dt.month() != last_bar.dt.month() || bar.dt.year() != last_bar.dt.year()
            }
            Freq::Quarter => {
                // 季线级别，检查是否跨季度
                let quarter1 = (bar.dt.month() - 1) / 3;
                let quarter2 = (last_bar.dt.month() - 1) / 3;
                quarter1 != quarter2 || bar.dt.year() != last_bar.dt.year()
            }
            Freq::Year => {
                // 年线级别，检查是否跨年
                bar.dt.year() != last_bar.dt.year()
            }
        }
    }

    /// 生成新的K线
    fn generate_new_bar(&self, current_bar: &RawBar, freq: &Freq) -> Result<RawBar, String> {
        let new_bar = RawBar::new(
            current_bar.symbol.clone(),
            current_bar.id,
            current_bar.dt,
            freq.clone(),
            current_bar.open,
            current_bar.close,
            current_bar.high,
            current_bar.low,
            current_bar.vol,
            current_bar.amount,
        );
        Ok(new_bar)
    }

    /// 更新当前K线
    fn update_current_bar(&self, last_bar: &mut RawBar, current_bar: &RawBar) {
        last_bar.high = last_bar.high.max(current_bar.high);
        last_bar.low = last_bar.low.min(current_bar.low);
        last_bar.close = current_bar.close;
        last_bar.vol += current_bar.vol;
        last_bar.amount += current_bar.amount;
    }

    /// 更新K线数据
    pub fn update(&mut self, bar: RawBar) -> Result<(), String> {
        // 验证频率设置
        self.validate_freqs()?;

        // 更新所有频率的K线
        let freqs = self.freqs.clone();
        for freq in &freqs {
            if let Some(freq_enum) = Freq::from_str(freq) {
                self.update_freq(&bar, &freq_enum)?;
            }
        }

        Ok(())
    }

    /// 获取指定频率的K线
    pub fn get_bars(&self, freq: &str) -> Option<&Vec<RawBar>> {
        self.freq_bars.get(freq)
    }

    /// 获取所有频率的K线
    pub fn get_all_bars(&self) -> &HashMap<String, Vec<RawBar>> {
        &self.freq_bars
    }
}

impl std::fmt::Display for BarGenerator {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "BarGenerator(base_freq={}, freqs={:?}, max_count={}, market={})",
            self.base_freq, self.freqs, self.max_count, self.market
        )
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use chrono::Utc;

    #[test]
    fn test_bar_generator() {
        let mut generator = BarGenerator::new(
            "1分钟".to_string(),
            vec!["1分钟".to_string(), "5分钟".to_string()],
            1000,
            "A股".to_string(),
        );

        let bar1 = RawBar::new(
            "000001.SZ".to_string(),
            1,
            Utc::now(),
            Freq::Min1,
            10.0,
            11.0,
            12.0,
            9.0,
            1000.0,
            11000.0,
        );

        let bar2 = RawBar::new(
            "000001.SZ".to_string(),
            2,
            Utc::now() + Duration::minutes(1),
            Freq::Min1,
            11.0,
            12.0,
            13.0,
            10.0,
            2000.0,
            24000.0,
        );

        generator.update(bar1).unwrap();
        generator.update(bar2).unwrap();

        assert!(generator.get_bars("1分钟").is_some());
        assert!(generator.get_bars("5分钟").is_some());
    }
} 